import numpy as np
import matplotlib.pyplot as plt
from pyswarm import pso
from tqdm import tqdm
import time

plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
plt.rc('font', size=10)
plt.rc('font', family='SimHei')

# 通用常数
interval = 1  # 投弹时间间隔
v_down = 3  # 云团下降速度
r_effective = 10  # 有效遮蔽半径
t_effective = 20  # 有效遮蔽时间
v_missile = 300  # 导弹速度
r_target = 7  # 目标半径
H_target = 10  # 目标高度
target = np.array([0, 200, 0])  # 目标底部圆心坐标
M1 = np.array([20000, 0, 2000])
M2 = np.array([19000, 600, 2100])
M3 = np.array([18000, -600, 1900])
FY1 = np.array([17800, 0, 1800])
FY2 = np.array([12000, 1400, 1400])
FY3 = np.array([6000, -3000, 700])
FY4 = np.array([11000, 2000, 1800])
FY5 = np.array([13000, -2000, 1300])
v_range = np.array([70, 140])  # 无人机速度范围
g = 9.8  # 重力加速度
target_point = np.array([
    [(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, 0],
    [-(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, 0],
    [(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, 0],
    [-(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, 0],
    [(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target],
    [-(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target],
    [(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target],
    [-(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target],
    [(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target / 2],
    [-(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target / 2],
    [(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target / 2],
    [-(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target / 2],
])
t_total = np.linalg.norm(M1) / 300  # 从0开始到导弹击中目标的总用时
precision = 0.01  # 时间精度

# 全局变量用于记录PSO优化过程
pso_history = {
    'iteration': [],
    'best_f': [],
    'mean_f': [],
    'params': [],
    'time': []
}


# 传入M1的初始坐标（也即可确定其飞行方向）和经过的时间，返回其在给定时间下的坐标
def get_M1_coordinate(coordinate_0, t, v=None):
    if not v:
        v = -coordinate_0 / np.linalg.norm(coordinate_0) * v_missile
    else:
        v = v / np.linalg.norm(v) * v_missile
    return coordinate_0 + v * t


# 传入第i架无人机的坐标和经过的时间，返回它投射的烟雾弹中心的坐标
def get_FY_coordinate(coordinate_0, t, i, v_1, t_1, t_2, theta):
    t_drop = t_1[i]
    t_explosion = t_2[i]
    if t < t_drop + t_explosion or t > t_drop + t_explosion + t_effective:
        return [None, None, None]
    else:
        v = np.array([np.cos(theta[i]), np.sin(theta[i]), 0]) * v_1[i]
        delta_t = t - t_drop - t_explosion
        coordinate = coordinate_0[i] + v * (t_drop + t_explosion)
        coordinate[2] -= 1 / 2 * g * (t_explosion ** 2)
        coordinate[2] -= delta_t * v_down
        return coordinate


def collision_detection(coordinate_missile, coordinate_0, t, v_1, t_1, t_2, theta):
    global target_point, r_effective
    count_collision = 0
    for target in target_point:
        for i in range(3):
            coordinate_bomb = get_FY_coordinate(coordinate_0, t, i, v_1, t_1, t_2, theta)
            if coordinate_bomb[0] is None:
                break
            l_1 = coordinate_bomb - target
            l_2 = coordinate_missile - target
            molecule = np.cross(l_1, l_2)
            denominator = np.linalg.norm(l_2)
            d = np.linalg.norm(molecule) / denominator
            if d < r_effective:
                missile_to_target = -coordinate_missile  # 由导弹指向目标
                missile_to_bomb = coordinate_bomb - coordinate_missile  # 由导弹指向烟幕弹
                cos = np.dot(missile_to_target, missile_to_bomb)
                if cos > 0 or np.linalg.norm(missile_to_bomb) < r_effective:  # 余弦值大于0，是锐角，或者导弹和烟幕弹之间的距离小于有效半径
                    count_collision += 1
                    break
    return count_collision


# 飞行方向、飞行速度、烟幕干扰弹投放点、烟幕干扰弹起爆点
def fun(theta, v_1, t_1, t_2):
    t_list = np.arange(0, t_total, precision)
    count = 0
    for t in t_list:
        if collision_detection(get_M1_coordinate(M1, t), [FY1, FY2, FY3], t, v_1, t_1, t_2, theta) == len(target_point):
            count += 1
    return count


# 定义变量范围
lb = [3.05, 70, 0, 0, 3, 70, 0, 0, 2.4, 70, 0, 0]
ub = [3.25, 140, 5, 5, 3.5, 140, 10, 10, 2.8, 140, 10, 10]


# 在函数外部定义一个列表来记录历史
history_best_x = []
history_best_f = []


# 定义目标函数
def objective(x):
    theta_0, v_1_0, t_1_0, t_2_0, theta_1, v_1_1, t_1_1, t_2_1, theta_2, v_1_2, t_1_2, t_2_2 = x
    theta = [theta_0, theta_1, theta_2]
    v_1 = [v_1_0, v_1_1, v_1_2]
    t_1 = [t_1_0, t_1_1, t_1_2]
    t_2 = [t_2_0, t_2_1, t_2_2]
    fitness = -fun(theta, v_1, t_1, t_2)  # 假设这是你原来的目标函数计算
    history_best_x.append(x.copy())
    history_best_f.append(fitness)
    return fitness  # PSO默认最小化，所以取负


# 自定义PSO回调函数
def pso_callback(x, f):
    global pso_history
    iteration = len(pso_history['iteration'])
    pso_history['iteration'].append(iteration)
    pso_history['best_f'].append(-f)  # 记录最佳适应度值
    pso_history['params'].append(x.copy())  # 记录参数
    pso_history['time'].append(time.time())

    # 更新进度条
    if hasattr(pso_callback, 'pbar'):
        pso_callback.pbar.update(1)
        pso_callback.pbar.set_postfix({'Best P': -f, 'Iteration': iteration})


# 绘制PSO优化过程
def plot_PSO_history():
    plt.figure(figsize=(15, 10))

    # 绘制适应度值变化
    plt.subplot(2, 3, 1)
    plt.plot(pso_history['iteration'], pso_history['best_f'], 'b-', label='最佳适应度')
    plt.xlabel('迭代次数')
    plt.ylabel('P值')
    plt.title('适应度值变化曲线')
    plt.legend()
    plt.grid(True)

    # 提取参数变化
    params = np.array(pso_history['params'])
    param_names = [
        'FY1飞行方向角', 'FY1速度', 'FY1投弹时间', 'FY1爆炸时间',
        'FY2飞行方向角', 'FY2速度', 'FY2投弹时间', 'FY2爆炸时间',
        'FY3飞行方向角', 'FY3速度', 'FY3投弹时间', 'FY3爆炸时间'
    ]

    # 绘制参数变化
    for i in range(4):
        plt.subplot(2, 3, i + 2)
        for j in range(3):
            idx = i + j * 4
            plt.plot(pso_history['iteration'], params[:, idx], label=param_names[idx])
        plt.xlabel('迭代次数')
        plt.ylabel('参数值')
        plt.title(f'参数组{i + 1}变化曲线')
        plt.legend()
        plt.grid(True)

    plt.tight_layout()
    plt.show()


# 运行PSO优化
print("开始粒子群优化...")
maxiter = 5
swarmsize = 3

# 创建进度条
with tqdm(total=maxiter, desc='优化进度') as pbar:
    pso_callback.pbar = pbar
    # 使用pso函数，并传入回调函数
    xopt, fopt = pso(objective, lb, ub, swarmsize=swarmsize, maxiter=maxiter,
                     debug=True)

# 提取最佳解
theta_0_b, v_1_0_b, t_1_0_b, t_2_0_b, theta_1_b, v_1_1_b, t_1_1_b, t_2_1_b, theta_2_b, v_1_2_b, t_1_2_b, t_2_2_b = xopt
theta_b = [theta_0_b, theta_1_b, theta_2_b]
v_1_b = [v_1_0_b, v_1_1_b, v_1_2_b]
t_1_b = [t_1_0_b, t_1_1_b, t_1_2_b]
t_2_b = [t_2_0_b, t_2_1_b, t_2_2_b]
Pb = -fopt  # 恢复目标值

print(f"最佳解: theta={theta_b}, v_1={v_1_b}, t_1={t_1_b}, t_2={t_2_b}, P={Pb}")

# # 绘制优化过程
# plot_PSO_history()

# 计算优化过程耗时
if len(pso_history['time']) > 0:
    optimization_time = pso_history['time'][-1] - pso_history['time'][0]
    print(f"优化过程耗时: {optimization_time:.2f} 秒")
    print(f"平均每代耗时: {optimization_time / maxiter:.2f} 秒")

# 绘制最佳解的参数分布
plt.figure(figsize=(12, 8))
param_names = [
    'FY1飞行方向角', 'FY1速度', 'FY1投弹时间', 'FY1爆炸时间',
    'FY2飞行方向角', 'FY2速度', 'FY2投弹时间', 'FY2爆炸时间',
    'FY3飞行方向角', 'FY3速度', 'FY3投弹时间', 'FY3爆炸时间'
]
plt.bar(param_names, xopt)
plt.xticks(rotation=45, ha='right')
plt.ylabel('参数值')
plt.title('最佳解参数分布')
plt.tight_layout()
plt.show()

# 添加收敛性分析
plt.figure(figsize=(10, 6))
best_f = np.array(pso_history['best_f'])
plt.semilogy(pso_history['iteration'], best_f, 'b-', label='最佳适应度')
plt.xlabel('迭代次数')
plt.ylabel('P值 (对数尺度)')
plt.title('收敛性分析')
plt.grid(True, which="both", ls="--")
plt.legend()
plt.tight_layout()
plt.show()

print(history_best_x, history_best_f)
